#!/usr/bin/env python
# -*- coding: utf-8 -*-

import time
from collections import Counter, defaultdict
from scapy.all import *

class PacketAnalyzer:
    """
    数据包分析类，负责对捕获的数据包进行统计和分析
    """
    
    def __init__(self):
        """
        初始化数据包分析器
        """
        # 初始化计数器
        self.total_packets = 0
        self.protocol_counter = Counter()
        self.ip_counter = Counter()
        self.port_counter = Counter()
        self.packet_size_distribution = defaultdict(int)
        self.packet_time_distribution = defaultdict(int)
        
        # 异常检测阈值
        self.port_scan_threshold = 50  # 提高端口扫描阈值到50个不同端口
        self.port_scan_time_window = 10  # 缩短时间窗口到10秒
        self.dos_threshold = 100       # 每秒向同一目标发送超过100个包视为DoS攻击
        
        # 异常检测状态
        self.potential_port_scans = defaultdict(list)  # 源IP -> [(时间戳, 目标端口, 目标IP)]
        self.potential_dos_attacks = defaultdict(list) # 目标IP -> [(时间戳, 源IP)]
        self.alerts = []
        
        # 协议映射
        self.protocol_map = {
            1: "ICMP",
            6: "TCP",
            17: "UDP",
            2: "IGMP"
        }
        
        # 是否启用端口扫描检测
        self.enable_port_scan_detection = False  # 默认关闭端口扫描检测
    
    def reset(self):
        """
        重置分析器状态
        """
        self.total_packets = 0
        self.protocol_counter.clear()
        self.ip_counter.clear()
        self.port_counter.clear()
        self.packet_size_distribution.clear()
        self.packet_time_distribution.clear()
        self.potential_port_scans.clear()
        self.potential_dos_attacks.clear()
        self.alerts = []
        # 不重置enable_port_scan_detection，保持用户的设置
    
    def analyze_packet(self, packet):
        """
        分析单个数据包，更新统计信息
        
        Args:
            packet: 可以是scapy原始数据包或者已解析的数据包信息字典
        """
        # 检查是否是已解析的数据包
        if isinstance(packet, dict):
            return self._analyze_parsed_packet(packet)
        
        # 如果是原始scapy数据包，先进行基本解析
        parsed_data = self._basic_parse_packet(packet)
        return self._analyze_parsed_packet(parsed_data)
    
    def _basic_parse_packet(self, packet):
        """
        对scapy数据包进行基本解析，提取关键字段
        
        Args:
            packet: scapy数据包
            
        Returns:
            dict: 包含基本信息的字典
        """
        try:
            parsed_data = {
                "length": len(packet),
                "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(packet.time)),
                "protocols": []
            }
            
            # 提取协议信息
            if Ether in packet:
                parsed_data["protocols"].append("以太网")
                parsed_data["src_mac"] = packet[Ether].src
                parsed_data["dst_mac"] = packet[Ether].dst
            
            if IP in packet:
                parsed_data["protocols"].append("IPv4")
                parsed_data["src_ip"] = packet[IP].src
                parsed_data["dst_ip"] = packet[IP].dst
                
                # 添加IP协议类型
                ip_proto = packet[IP].proto
                if ip_proto in self.protocol_map:
                    proto_name = self.protocol_map[ip_proto]
                    if proto_name not in parsed_data["protocols"]:
                        parsed_data["protocols"].append(proto_name)
            
            elif IPv6 in packet:
                parsed_data["protocols"].append("IPv6")
                parsed_data["src_ip"] = packet[IPv6].src
                parsed_data["dst_ip"] = packet[IPv6].dst
                
                # 添加下一头部类型
                nh = packet[IPv6].nh
                if nh in self.protocol_map:
                    proto_name = self.protocol_map[nh]
                    if proto_name not in parsed_data["protocols"]:
                        parsed_data["protocols"].append(proto_name)
            
            elif ARP in packet:
                parsed_data["protocols"].append("ARP")
                parsed_data["src_ip"] = packet[ARP].psrc
                parsed_data["dst_ip"] = packet[ARP].pdst
            
            # 提取传输层信息
            if TCP in packet:
                if "TCP" not in parsed_data["protocols"]:
                    parsed_data["protocols"].append("TCP")
                parsed_data["src_port"] = str(packet[TCP].sport)
                parsed_data["dst_port"] = str(packet[TCP].dport)
                
                # 检测应用层协议
                if packet[TCP].sport == 80 or packet[TCP].dport == 80:
                    parsed_data["protocols"].append("HTTP")
                elif packet[TCP].sport == 443 or packet[TCP].dport == 443:
                    parsed_data["protocols"].append("HTTPS")
                elif packet[TCP].sport == 21 or packet[TCP].dport == 21:
                    parsed_data["protocols"].append("FTP")
                elif packet[TCP].sport == 22 or packet[TCP].dport == 22:
                    parsed_data["protocols"].append("SSH")
                elif packet[TCP].sport == 25 or packet[TCP].dport == 25:
                    parsed_data["protocols"].append("SMTP")
            
            elif UDP in packet:
                if "UDP" not in parsed_data["protocols"]:
                    parsed_data["protocols"].append("UDP")
                parsed_data["src_port"] = str(packet[UDP].sport)
                parsed_data["dst_port"] = str(packet[UDP].dport)
                
                # 检测应用层协议
                if packet[UDP].sport == 53 or packet[UDP].dport == 53:
                    parsed_data["protocols"].append("DNS")
                elif packet[UDP].sport == 67 or packet[UDP].dport == 67 or packet[UDP].sport == 68 or packet[UDP].dport == 68:
                    parsed_data["protocols"].append("DHCP")
            
            elif ICMP in packet:
                if "ICMP" not in parsed_data["protocols"]:
                    parsed_data["protocols"].append("ICMP")
            
            # 确定顶层协议
            if parsed_data["protocols"]:
                parsed_data["protocol"] = parsed_data["protocols"][-1]
            else:
                parsed_data["protocol"] = "未知"
            
            # 构建源和目标
            src = parsed_data.get("src_ip", "未知")
            dst = parsed_data.get("dst_ip", "未知")
            
            if "src_port" in parsed_data:
                parsed_data["src"] = f"{src}:{parsed_data['src_port']}"
            else:
                parsed_data["src"] = src
                
            if "dst_port" in parsed_data:
                parsed_data["dst"] = f"{dst}:{parsed_data['dst_port']}"
            else:
                parsed_data["dst"] = dst
            
            return parsed_data
            
        except Exception as e:
            # 如果解析失败，返回最小信息集
            return {
                "protocol": "未知",
                "src": "未知",
                "dst": "未知",
                "length": len(packet) if packet else 0,
                "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
            }
    
    def _analyze_parsed_packet(self, parsed_packet):
        """
        分析已经解析好的数据包
        
        Args:
            parsed_packet: 解析后的数据包信息字典
        """
        # 增加总数据包计数
        self.total_packets += 1
        
        # 更新协议统计
        if "protocol" in parsed_packet:
            self.protocol_counter[parsed_packet["protocol"]] += 1
        
        # 更新IP地址统计
        if "src" in parsed_packet and parsed_packet["src"] != "未知":
            src_ip = parsed_packet["src"].split(":")[0]  # 去除端口部分
            self.ip_counter[f"源: {src_ip}"] += 1
            
        if "dst" in parsed_packet and parsed_packet["dst"] != "未知":
            dst_ip = parsed_packet["dst"].split(":")[0]  # 去除端口部分
            self.ip_counter[f"目标: {dst_ip}"] += 1
            
        # 更新端口统计
        if "src_port" in parsed_packet and parsed_packet["src_port"]:
            self.port_counter[f"源端口: {parsed_packet['src_port']}"] += 1
            
        if "dst_port" in parsed_packet and parsed_packet["dst_port"]:
            self.port_counter[f"目标端口: {parsed_packet['dst_port']}"] += 1
            
        # 更新数据包大小分布
        if "length" in parsed_packet and parsed_packet["length"] != "N/A":
            try:
                length = int(parsed_packet["length"])
                # 按照大小范围分组
                if length < 64:
                    size_range = "<64"
                elif length < 128:
                    size_range = "64-127"
                elif length < 256:
                    size_range = "128-255"
                elif length < 512:
                    size_range = "256-511"
                elif length < 1024:
                    size_range = "512-1023"
                elif length < 1500:
                    size_range = "1024-1499"
                else:
                    size_range = ">=1500"
                    
                self.packet_size_distribution[size_range] += 1
            except:
                pass
                
        # 更新时间分布
        if "time" in parsed_packet:
            try:
                # 提取小时和分钟
                time_parts = parsed_packet["time"].split()
                if len(time_parts) > 1:
                    time_str = time_parts[1][:5]  # 取HH:MM部分
                    self.packet_time_distribution[time_str] += 1
            except:
                pass
                
        # 进行异常检测
        self._check_for_anomalies(parsed_packet)
    
    def get_protocol_stats(self):
        """
        获取协议统计信息
        
        Returns:
            list: 协议统计信息，格式为[(协议名称, 数量), ...]
        """
        return self.protocol_counter.most_common()
    
    def get_ip_stats(self, top_n=10):
        """
        获取IP地址统计信息
        
        Args:
            top_n: 返回前N个最活跃的IP地址
            
        Returns:
            list: IP统计信息，格式为[(IP地址, 数量), ...]
        """
        return self.ip_counter.most_common(top_n)
    
    def get_port_stats(self, top_n=10):
        """
        获取端口统计信息
        
        Args:
            top_n: 返回前N个最活跃的端口
            
        Returns:
            list: 端口统计信息，格式为[(端口, 数量), ...]
        """
        return self.port_counter.most_common(top_n)
    
    def get_packet_size_distribution(self):
        """
        获取数据包大小分布
        
        Returns:
            dict: 数据包大小分布，格式为{大小范围: 数量, ...}
        """
        # 确保所有范围都存在，即使计数为0
        size_ranges = ["<64", "64-127", "128-255", "256-511", "512-1023", "1024-1499", ">=1500"]
        result = {size_range: self.packet_size_distribution.get(size_range, 0) for size_range in size_ranges}
        return result
    
    def get_packet_time_distribution(self):
        """
        获取数据包时间分布
        
        Returns:
            dict: 数据包时间分布，格式为{时间: 数量, ...}
        """
        # 按时间排序
        sorted_times = sorted(self.packet_time_distribution.items())
        return dict(sorted_times)
    
    def get_packet_rate(self, time_window=60):
        """
        计算最近时间窗口内的数据包速率
        
        Args:
            time_window: 时间窗口大小，单位为秒
            
        Returns:
            float: 每秒数据包数量
        """
        # 这里简单地返回总包数除以时间窗口
        # 实际上可以实现更复杂的速率计算，例如使用滑动窗口
        if self.total_packets == 0:
            return 0
        return self.total_packets / time_window
    
    def get_alerts(self):
        """
        获取检测到的异常警报
        
        Returns:
            list: 警报信息列表
        """
        return self.alerts
    
    def _check_for_anomalies(self, packet):
        """
        检查数据包是否存在异常行为
        
        Args:
            packet: 解析后的数据包信息
        """
        try:
            # 只有包含源IP、目标IP和端口信息的包才能检测异常
            if "src" not in packet or "dst" not in packet or packet["src"] == "未知" or packet["dst"] == "未知":
                return
                
            src_ip = packet["src"].split(":")[0]
            dst_ip = packet["dst"].split(":")[0]
            dst_port = packet.get("dst_port", "")
            timestamp = time.time()  # 当前时间戳
            
            # 检查端口扫描（只有启用时才检测）
            if self.enable_port_scan_detection and dst_port:
                # 记录目标IP和端口的组合，只有向同一IP的不同端口才算扫描
                self.potential_port_scans[src_ip].append((timestamp, dst_port, dst_ip))
                
                # 清理时间窗口外的记录
                self.potential_port_scans[src_ip] = [
                    (ts, port, ip) for ts, port, ip in self.potential_port_scans[src_ip]
                    if timestamp - ts <= self.port_scan_time_window
                ]
                
                # 按目标IP分组，统计每个目标IP上访问的不同端口数量
                ip_port_map = defaultdict(set)
                for _, port, ip in self.potential_port_scans[src_ip]:
                    ip_port_map[ip].add(port)
                
                # 检查是否有任何单个目标IP的不同端口数量超过阈值
                for target_ip, ports in ip_port_map.items():
                    if len(ports) >= self.port_scan_threshold:
                        alert = f"可能的端口扫描行为: 源IP {src_ip} 在{self.port_scan_time_window}秒内向目标IP {target_ip} 访问了 {len(ports)} 个不同端口"
                        if alert not in self.alerts:
                            self.alerts.append(alert)
            
            # 检查DoS攻击
            self.potential_dos_attacks[dst_ip].append((timestamp, src_ip))
            # 清理1秒前的记录
            self.potential_dos_attacks[dst_ip] = [
                (ts, ip) for ts, ip in self.potential_dos_attacks[dst_ip]
                if timestamp - ts <= 1
            ]
            
            # 检查是否超过阈值
            if len(self.potential_dos_attacks[dst_ip]) > self.dos_threshold:
                # 统计源IP分布，检查是否分布式攻击
                dos_sources = Counter([ip for _, ip in self.potential_dos_attacks[dst_ip]])
                if len(dos_sources) < 3:
                    # 单源攻击
                    top_source = dos_sources.most_common(1)[0][0]
                    alert = f"可能的DoS攻击: 目标IP {dst_ip} 在1秒内收到了 {len(self.potential_dos_attacks[dst_ip])} 个数据包，主要来自 {top_source}"
                else:
                    # 分布式攻击
                    alert = f"可能的DDoS攻击: 目标IP {dst_ip} 在1秒内收到了 {len(self.potential_dos_attacks[dst_ip])} 个数据包，来自 {len(dos_sources)} 个不同源IP"
                    
                if alert not in self.alerts:
                    self.alerts.append(alert)
        except Exception as e:
            # 异常检测错误不应该影响主流程
            pass
    
    def get_summary(self):
        """
        获取数据包分析摘要
        
        Returns:
            dict: 包含各类统计信息的字典
        """
        return {
            "total_packets": self.total_packets,
            "protocols": self.get_protocol_stats(),
            "top_ips": self.get_ip_stats(5),
            "top_ports": self.get_port_stats(5),
            "size_distribution": self.get_packet_size_distribution(),
            "alerts": self.get_alerts()
        } 